# src/data_handler.py

import os
import random
import requests
from llama_index.core import SimpleDirectoryReader, Document
from llama_index.llms.openai import OpenAI
from llama_index.core.evaluation import DatasetGenerator
from typing import List

from . import config


def download_data():
    """下载 PDF 数据文件，如果不存在的话。"""
    if not os.path.exists(config.DATA_DIR):
        os.makedirs(config.DATA_DIR)

    if not os.path.exists(config.PDF_PATH):
        print(f"正在下载数据文件从 {config.PDF_URL}...")
        response = requests.get(config.PDF_URL)
        response.raise_for_status()  # 确保请求成功
        with open(config.PDF_PATH, "wb") as f:
            f.write(response.content)
        print(f"文件已保存至 {config.PDF_PATH}")
    else:
        print(f"数据文件 {config.PDF_PATH} 已存在。")


def load_and_split_documents() -> (List[Document], List[Document]):
    """加载 PDF 文档并分割为训练集和评估集。"""
    print("正在加载和分割文档...")
    documents = SimpleDirectoryReader(input_files=[config.PDF_PATH]).load_data()

    # 打乱文档顺序以确保随机性
    random.seed(config.RANDOM_SEED)
    random.shuffle(documents)

    train_docs = documents[:config.TRAIN_DOCS_LIMIT]
    eval_docs = documents[config.EVAL_DOCS_START:]

    print(f"文档加载完成。训练集: {len(train_docs)} 页, 评估集: {len(eval_docs)} 页。")
    return train_docs, eval_docs


def generate_and_save_questions(documents: List[Document], output_path: str, num_questions: int):
    """使用给定的文档生成问题并保存到文件。"""
    if os.path.exists(output_path):
        print(f"问题文件 {output_path} 已存在，跳过生成。")
        return

    print(f"正在为 {os.path.basename(output_path)} 生成 {num_questions} 个问题...")

    llm = OpenAI(model=config.BASE_MODEL, temperature=config.BASE_MODEL_TEMP)

    dataset_generator = DatasetGenerator.from_documents(
        documents,
        question_gen_query=config.QUESTION_GEN_QUERY,
        llm=llm,
    )

    questions = dataset_generator.generate_questions_from_nodes(num=num_questions)

    with open(output_path, "w") as f:
        for question in questions:
            f.write(question.strip() + "\n")

    print(f"成功生成并保存 {len(questions)} 个问题到 {output_path}")


def load_questions(file_path: str) -> List[str]:
    """从文件中加载问题列表。"""
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"问题文件 {file_path} 未找到。请先生成问题。")

    with open(file_path, "r") as f:
        return [line.strip() for line in f]